#include "include/types.h"
#include "include/param.h"
#include "include/memlayout.h"
#include "include/riscv.h"
#include "include/spinlock.h"
#include "include/sleeplock.h"
#include "include/proc.h"
#include "include/defs.h"
#include "include/elf.h"
#include "include/fat32.h"
#include "include/binformat.h"
#include "include/debug.h"

const char *env[] = {"SHELL=shell",
                     "PWD=/",
                     "HOME=/",
                     "USER=root",
                     "MOTD_SHOWN=pam",
                     "LANG=C.UTF-8",
                     "INVOCATION_ID=e9500a871cf044d9886a157f53826684",
                     "TERM=vt220",
                     "SHLVL=2",
                     "JOURNAL_STREAM=8:9265",
                     "PATH=/bin",
                     "OLDPWD=/root",
                     "_=busybox",
                     0};


uint64 CreateUserStack(struct BinProgram *bin_program, struct elfhdr *elf)
{
  int index = bin_program->argc + bin_program->envc + 2;

  uint64 filename = CopyString("/bin/busybox", bin_program);
#define NEW_AUX_ENT(id, val)                                                   \
  do {                                                                         \
    bin_program->ustack[index++] = id;                                         \
    bin_program->ustack[index++] = val;                                        \
  } while (0)

  NEW_AUX_ENT(0x28, 0);
  NEW_AUX_ENT(0x29, 0);
  NEW_AUX_ENT(0x2a, 0);
  NEW_AUX_ENT(0x2b, 0);
  NEW_AUX_ENT(0x2c, 0);
  NEW_AUX_ENT(0x2d, 0);
  NEW_AUX_ENT(AT_PHDR, elf->phoff);               // 3
  NEW_AUX_ENT(AT_PHENT, sizeof(struct proghdr));  // 4
  NEW_AUX_ENT(AT_PHNUM, elf->phnum);              // 5
  NEW_AUX_ENT(AT_PAGESZ, PGSIZE);                 // 6
  NEW_AUX_ENT(AT_BASE, 0);                        // 7
  NEW_AUX_ENT(AT_FLAGS, 0);                       // 8
  NEW_AUX_ENT(AT_ENTRY, elf->entry);              // 9
  NEW_AUX_ENT(AT_UID, 0);                         // 11
  NEW_AUX_ENT(AT_EUID, 0);                        // 12
  NEW_AUX_ENT(AT_GID, 0);                         // 13
  NEW_AUX_ENT(AT_EGID, 0);                        // 14
  NEW_AUX_ENT(AT_HWCAP, 0x112d);                  // 16
  NEW_AUX_ENT(AT_CLKTCK, 64);                     // 17
  NEW_AUX_ENT(AT_EXECFN, filename);               // 31
  NEW_AUX_ENT(0, 0);
#undef NEW_AUX_ENT
  bin_program->sp -= sizeof(uint64) * index;
  if (copyout(bin_program->pagetable, bin_program->sp,
              (char *)bin_program->ustack, sizeof(uint64) * index)) {
    return -1;
  }
  uint64 argc = bin_program->argc;
  bin_program->sp -= sizeof(uint64);
  __debug_info("exec", "argc stack pointer = 0x%x\n", bin_program->sp);
  copyout(bin_program->pagetable, bin_program->sp, (char *)&argc, sizeof(uint64));
  return 0;
}
static int
loadseg(pagetable_t pagetable, uint64 va, struct dirent *ep, uint offset, uint sz)
{
  uint i, n;
  uint64 pa;
  for (i = 0; i < sz; i += n) {
    pa = walkaddr(pagetable, va + i);
    if (pa == 0) {
      printf("load seg error, address should exist\n");
    }
    pa += (va + i) % PGSIZE;
    n = PGSIZE - pa % PGSIZE;
    if ((sz - i) < n) {
      n = sz - i;
    }
    if (eread(ep, 0, pa, offset + i, n) != n) {
      printf("read error, bytes should equal n\n");
      return -1;
    }
  }
  return 0;
}

int exec(char *path, char **argv)
{
  char *s, *last;
  int i, off;
  uint64 sz = 0, sp, ustack[MAXARG+1], stackbase, oldsz;
  struct BinProgram  bin_prog;
  struct elfhdr elf;
  struct dirent *ep;
  struct proghdr ph;
  pagetable_t pagetable = 0, oldpagetable;
  struct proc *p = myproc();
  if((ep = ename(path)) == 0) {
    return -1;
  }
  elock(ep);
  // printf("argv[0] = %s\n", argv[0]);
  // printf("argv[1] = %s\n", argv[1]);
  // printf("argv[2] = %s\n", argv[2]);
  // printf("argv[3] = %s\n", argv[3]);
  // Check ELF header
  if(eread(ep, 0, (uint64) &elf, 0, sizeof(elf)) != sizeof(elf))
    goto bad;
  if(elf.magic != ELF_MAGIC)
    goto bad;

  if((pagetable = proc_pagetable(p)) == 0)
    goto bad;

  // Load program into memory.
  for(i=0, off=elf.phoff; i<elf.phnum; i++, off+=sizeof(ph)){
    if(eread(ep, 0, (uint64)&ph, off, sizeof(ph)) != sizeof(ph)){
      goto bad;
    }
    if(ph.type != ELF_PROG_LOAD) // only load type segment need map 
      continue;
    if(ph.memsz < ph.filesz) // memsz must > filesz, we will alloc memsz & map filesz
      goto bad;
    if(ph.vaddr + ph.memsz < ph.vaddr)
      goto bad;
    uint64 sz1;
    if((sz1 = uvmalloc(pagetable, sz, ph.vaddr + ph.memsz)) == 0)
      goto bad;
    sz = sz1;
    if(loadseg(pagetable, ph.vaddr, ep, ph.off, ph.filesz) < 0)
      goto bad;
  }
  eunlock(ep);
  eput(ep);
  ep = 0;

  p = myproc();
  oldsz = p->sz;

  // Allocate 4 pages at the next page boundary. Use 2,3,4 as the user stack.
  sz = PGROUNDUP(sz);
  uint64 sz1;
  if((sz1 = uvmalloc(pagetable, sz, sz + 4*PGSIZE)) == 0)
    goto bad;
  sz = sz1;
  // Make the first page invaild
  uvmclear(pagetable, sz-3*PGSIZE); 
  sp = sz;
  stackbase = sp - PGSIZE;
  sp -= sizeof(uint64);

  // Fork from startOS by huiyu
  bin_prog.sp = sp;
  bin_prog.stackbase = stackbase;
  bin_prog.pagetable = pagetable;
  bin_prog.ustack = ustack;

  // push environment strings
  bin_prog.envc = CopyString2Stack((char **)env, &bin_prog);
  __debug_info("exec", "envc = %d\n", bin_prog.envc);
  __debug_info("exec", "env string stack pointer = 0x%x\n", bin_prog.sp);
  // push argument strings
  bin_prog.argc = CopyString2Stack(argv, &bin_prog);
  __debug_info("exec", "argc = %d\n", bin_prog.argc);
  __debug_info("exec", "arg string stack pointer = 0x%x\n", bin_prog.sp);
  // push padding
  bin_prog.sp -= bin_prog.sp % 16;
  // push string identify platform
  char k_platform [] = "RISC-V64";
  bin_prog.sp -= strlen(k_platform);
  if(copyout(pagetable, bin_prog.sp, k_platform, strlen(k_platform)) < 0)
    goto bad;
  __debug_info("exec", "platfrom string stack pointer = 0x%x\n", bin_prog.sp);
  // push rand bytes
  bin_prog.sp -= 16;
  char rand_bytes [] = "012345678912345"; 
  if(copyout(pagetable, bin_prog.sp, rand_bytes, strlen(rand_bytes) + 1) < 0)
    goto bad;
  __debug_info("exec", "rand bytes stack pointer = 0x%x\n", bin_prog.sp);
  // push padding
  bin_prog.sp -= bin_prog.sp % 16;
  // push argc, argv[], envp[], auxv[]
  CreateUserStack(&bin_prog, &elf);

  sp = bin_prog.sp;

  // arguments to user main(argc, argv)
  p->trapframe->a1 = sp;
  // Save program name for debugging.
  for(last=s=path; *s; s++)
    if(*s == '/')
      last = s+1;
  safestrcpy(p->name, last, sizeof(p->name));
  
  // Commit to the user image.
  oldpagetable = p->pagetable;
  p->pagetable = pagetable;
  p->sz = sz;
  p->trapframe->epc = elf.entry;  // initial program counter = main
  p->trapframe->sp = sp;          // initial stack pointer
  proc_freepagetable(oldpagetable, oldsz);
  
  return 0; // this ends up in a0

 bad:
  if(pagetable)
    proc_freepagetable(pagetable, sz);
  if(ep){
    eunlock(ep);
    eput(ep);
  }
  return -1;
}
